Simulating A Pendulum using SciPy

WIP (need to add documentation)

import numpy as np
import scipy.integrate as spi
import matplotlib.pyplot as plt
import matplotlib.animation as anim
from IPython.display import HTML
%matplotlib notebook

The Pendulum Equation

Second order differential equation:

$$ \ddot{\theta} + \frac{b}{m} \dot{\theta} + \frac{g}{L} sin(\theta) = 0 $$

Expand it to a vectorized form (using $\omega = \dot{\theta}$ as angular velocity)

$$ \begin{bmatrix} \dot{\theta} \\ \dot{\omega} \end{bmatrix} = \begin{bmatrix} \omega \\ -\frac{g}{L} sin(\theta) -\frac{b}{m} \omega \end{bmatrix} $$

def pendulum(t, Th, g=9.81, L=1, m=3, b=0.5):
    """
    Pendulum differential equation
    """
    th, om = Th
    omd =  - (b/m) * om - (g/L) * np.sin(th)
    out = np.array([ om, omd ])
    return out

Generate Phase Plot of Equation

thr, omr = np.meshgrid(
    np.linspace(-2*np.pi, 2*np.pi, 30),
    np.linspace(-2*np.pi, 2*np.pi, 30)
)
thv, omv = pendulum(0.0, [thr, omr])
plt.figure()
plt.quiver(thr, omr, thv, omv, np.hypot(thv, omv))
plt.xlabel('$\\theta$ (radians)')
plt.ylabel('$\omega$ (rads/sec)')
plt.title('Phase Plot of Pendulum Equation')
plt.axis([ -2*np.pi, 2*np.pi, -2*np.pi, 2*np.pi ])
plt.show()

Generate Phase Flows using scipy.integrate

# Equation meta parameters
L = 10
g = 9.81
m = 3
b = 0.5

# Time span
t = (0, 20)
max_dt = 0.01

# Initial parameters
Th0s = np.radians([
    [  30,   0 ],
    [  45,   0 ],
    [  60,   0 ],
    [  90,   0 ],
    [ 150,   0 ],
    [ 190,   0 ],
    [   0,  30 ],
    [   0,  45 ],
    [   0,  60 ],
    [   0,  90 ],
    [   0, 150 ],
    [   0, 180 ]
])

# Do all the integrations
Ths = [
    spi.solve_ivp(
        pendulum, 
        t, Th0, 
        max_step=max_dt, 
        args=(L, g, m, b)).y
    for Th0 in Th0s ]

# Plot them all!
plt.figure()
for Th in Ths:
    plt.plot(Th[0], Th[1])
plt.axis([-2*np.pi, 2*np.pi, -2*np.pi, 2*np.pi])
plt.xlabel('$\\theta$ (radians)')
plt.ylabel('$\\omega$ (rads/sec)')
plt.title('Flow Plot of Pendulum')
plt.show()

Simulate Multiple Pendulums

# Initial parameters
Th0 = np.radians([
    # Initial angles
    [ 30,  60,  90, 190,   0,   0,   0,   0],
    # Initial angular velocities
    [  0,   0,   0,   0,  30,  60,  90, 180]
])

# Equation meta parameters
L = 1
g = 9.81
m = 3
b = 0.5

# Figure parameters
interval = 12
framerate = 60
size = 1.5

# Turning interactive plot off temporarily
plt.ioff()

# Plot dimensions
dims = 8, 6
ylim = -size, size
xsize = size * dims[0] / dims[1]
xlim = -xsize, xsize

# Create figure
fig = plt.figure(figsize=dims)
ax = plt.axes(xlim=xlim, ylim=ylim)
lines = [
    ax.plot([], [], lw=1, marker='o')[0]
    for i in range(Th0.shape[1])
]
plt.xticks([])
plt.yticks([])
plt.title('Pendulum Simulation!')

def init():
    """
    Initialization function
    """
    for line in lines:
        line.set_data([], [])
    return lines

t = 0.0
Th = Th0
def animate(i):
    """
    Animation function
    """
    global t
    global Th
    dt = 1/framerate
    t += dt
    Th += pendulum(t, Th, g, L, m, b)*dt
    th = Th[0,:]
    x = L*np.sin(th)
    y = -L*np.cos(th)
    for i,line in enumerate(lines):
        line.set_data([0, x[i]],[0, y[i]])
    return lines

# Generate animation
frames = interval*framerate
anm = anim.FuncAnimation(
    fig, animate,
    init_func=init,
    interval=interval,
    frames=frames,
    blit=True)

# Close plot
plt.close(fig)
plt.ion()

# Render animation
HTML(anm.to_jshtml())

Once Loop Reflect